require 'config'
local match = string.match
local ngxmatch=ngx.re.match
local unescape=ngx.unescape_uri
local get_headers = ngx.req.get_headers
local optionIsOn = function (options) return options == "on" and true or false end
logpath = logdir 
rulepath = RulePath

Attack_Log = optionIsOn(Attack_Log_Module)

Ip_White_List_Module = optionIsOn(Ip_White_List_Module) -- pass waf
Host_White_List_Module = optionIsOn(Host_White_List_Module) -- pass waf
Ip_White_List_limit_Module = optionIsOn(Ip_White_List_limit_Module) -- wait random seconds
Ip_Black_List_Module = optionIsOn(Ip_Black_List_Module)  -- direct return 403
User_Agent_Black_List_Module = optionIsOn(User_Agent_Black_List_Module) -- direct return 403
Url_Black_List_Module = optionIsOn(Url_Black_List_Module) -- no rate limit ,direct return 403
Http_Method_Module = optionIsOn(Http_Method_Module)

Url_List_CC_Module = optionIsOn(Url_List_CC_Module) -- rate limit ,above rate return 403
Base_Ip_Url_CC_Module = optionIsOn(Base_Ip_Url_CC_Module)    
Base_Ip_CC_Module = optionIsOn(Base_Ip_CC_Module)
Redirect = optionIsOn(Redirect)
Mime_Type_Base_Ip_CC_Module = optionIsOn(Mime_Type_Base_Ip_CC_Module)

-- some common functions define black
function GetCurrentMimeType()
    local pre_mime_type = ngx.req.get_headers()['Accept']
    if pre_mime_type then
        local from, to, err = ngx.re.find(pre_mime_type, ",", "o")
        if from then
            local m,err = ngx.re.match(pre_mime_type, '^(.+?),')
            mime_type = m[1]
        else
            if pre_mime_type == "*/*" then
                local regex = "\\.(\\w+?)$"
                local m, err = ngx.re.match(ngx.var.uri, regex, "o")
                if m then
                    mime_type = m[1]
                else
                    mime_type = "other"
                end
            else
                mime_type = pre_mime_type
            end
        end
    else
        local regex = "\\.(\\w+?)$"
        local m, err = ngx.re.match(ngx.var.uri, regex, "o")
        if m then
            mime_type = m[1]
        else
            mime_type = "other"
        end
    end
    return mime_type
end

function getClientIp()
        -- IP = ngx.req.get_headers()["X-Real-IP"]
        IP = ngx.var.http_x_forwarded_for
        -- local regex = [[\\d+?.\d+?.\d+?.\d+?]] 
        if IP == nil then
                IP  = ngx.var.remote_addr 
        else
                IP = ngx.re.match(IP,"[0-9]+.[0-9]+.[0-9]+.[0-9]+")[0]
        end
        if IP == nil then
                IP  = "unknown"
        end
        return IP
end
function write_log(logfile,msg)
    local fd = io.open(logfile,"ab")
    if fd == nil then return end
    fd:write(msg)
    fd:flush()
    fd:close()
end
------------------------------------
function log(method,url,data,ruletag)
    if Attack_Log then
        local realIp = getClientIp()
        local ua = ngx.var.http_user_agent
        local servername=ngx.var.server_name
        local time=ngx.localtime()
        if ua  then
            line = realIp.." ["..time.."] \""..method.." "..servername..url.."\" \""..data.."\"  \""..ua.."\" \""..ruletag.."\"\n"
        else
            line = realIp.." ["..time.."] \""..method.." "..servername..url.."\" \""..data.."\" - \""..ruletag.."\"\n"
        end
        local filename = logpath..'/'..ngx.today().."_sec.log"
        write_log(filename,line)
    end
end
------------------------------------read rules function-------------------------------------------------------------------
function read_rule(var)
    file = io.open(rulepath..'/'..var,"r")
    if file==nil then
        return
    end
    t = {}
    for line in file:lines() do
        table.insert(t,line)
    end
    file:close()
    return(t)
end

function read_file_to_regex_rule(var)
    file = io.open(rulepath..'/'..var,"r")
    if file==nil then
        return
    end
    local t = {}
    for line in file:lines() do
        table.insert(t,line)
    end
    file:close()
    return(table.concat(t,"|"))
end

ip_white_list_rules = read_file_to_regex_rule('ip_white_list') 
ip_white_list_limit_rules = read_file_to_regex_rule('ip_white_list_limit_rules')
ip_black_list_rules = read_file_to_regex_rule('ip_black_list') -- white_list's priority > black_list's priority
url_black_list_rules = read_rule('url_black_list') -- if ngx.var.request_uri in list return 403;Note url_black_list's priority > url_cc_deny's priority 
url_cc_deny_rules = read_file_to_regex_rule('url_cc_deny') -- if ngx.var.request_uri in list  request into CCdeny limit rate
user_agent_deny_rules = read_rule('user_agent') -- ngx.var.http_user_agent in list return 403 
host_white_rules = read_file_to_regex_rule('host_white')



function url_black_list_module()
    if Url_Black_List_Module then
        if url_black_list_rules ~="" and ngx.re.find(ngx.var.request_uri,url_black_list_rules,"imjo") then
            log('BLOCK_URL',ngx.var.request_uri,"-",url_black_list_rules)
            ngx.exit(403)
            return true
        end
    end
    return false
end

function http_method_module()
    local method=ngx.req.get_method()
    if Http_Method_Module then
        if not ngx.re.find(method,Http_Method_Module_Rule,"imjo") then
            log('Invalid_method',ngx.var.request_uri,"-",Http_Method_Module_Rule)
            ngx.exit(403)
            return true
        end
    end
    return false
end

function user_agent_module()
    local ua = ngx.var.http_user_agent
    if ua ~= nil and User_Agent_Black_List_Module then
        for _,rule in pairs(user_agent_deny_rules) do
            if rule ~="" and ngxmatch(ua,rule,"imjo") then
                log('UA_BLOCKLIST',ngx.var.request_uri,"-",rule)
                ngx.exit(403)
            return true
            end
        end
    end
    return false
end

function host_white_list()
    local host = ngx.var.host
    if host ~= nil and Host_White_List_Module then
        if host_white_rules ~="" and ngx.re.find(host,host_white_rules,"imjo") then
            -- log('host_white_list',ngx.var.request_uri,"-",host_white_rules)
        	return true
        end
    end
    return false
end

function ip_white_list_module()
    if Ip_White_List_Module then
        if ip_white_list_rules ~="" and ngx.re.find(getClientIp(),ip_white_list_rules,"imjo") then
            -- log('ip_white_list',ngx.var.request_uri,"-",ip_white_list_rules)
            return true
        end
    end
    return false
end

function ip_black_list_module()
     if Ip_Black_List_Module then
         if ip_black_list_rules ~="" and ngx.re.find(getClientIp(),ip_black_list_rules,"imjo") then
             log('GET',ngx.var.request_uri,"-","BlackIp")
             ngx.exit(403)
             return true
         end
     end
         return false
end

function base_ip_url_cc_module()
    -- in this function use source_ip and request_uri as token.It is global,but you can work with nginx location{}.
    if Base_Ip_Url_CC_Module then
        local uri=ngx.var.request_uri
        CCcount=tonumber(string.match(Base_Ip_Url_CC_Rate,'(.*)/'))
        CCseconds=tonumber(string.match(Base_Ip_Url_CC_Rate,'/(.*)'))
        local token = getClientIp()..uri -- fingerprints 
        local limit = ngx.shared.limit
        local req,_=limit:get(token)
        if req then
            if req > CCcount then
                 log('GET',ngx.var.request_uri,"-","Base_Ip_Url_CC")
                 ngx.exit(403)
                return true
            else
                 limit:incr(token,1)
            end
        else
            limit:set(token,1,CCseconds)
        end
    end
    return false
end

function base_ip_cc_module()
    -- in this function just use ip as token,if use it in http{},suggest set is large
    if Base_Ip_CC_Module then
        CCcount = tonumber(string.match(Base_Ip_CC_Rate,'(.*)/'))
        CCseconds = tonumber(string.match(Base_Ip_CC_Rate,'/(.*)'))
        local token = getClientIp() -- fingerprints 
        local limit = ngx.shared.limit
        local req,_ = limit:get(token)
        if req then
            if req > CCcount then
                log('GET',ngx.var.request_uri,"-","Base_Ip_CC")
                ngx.exit(403)
                return true
            else
                limit:incr(token,1)
            end
        else
            limit:set(token,1,CCseconds)
        end
    end
    return false
end

function base_url_list_cc_module()
    -- you can use it in http{},if request_uri not in black list ,just pass.
    if Url_List_CC_Module then
        if url_cc_deny_rules ~="" and ngx.re.find(ngx.var.uri,url_cc_deny_rules,"imjo") then
            CCcount = tonumber(string.match(Url_List_CC_Rate,'(.*)/'))
            CCseconds = tonumber(string.match(Url_List_CC_Rate,'/(.*)'))
            local token = getClientIp().."blacklisturl" -- fingerprints 
            local limit = ngx.shared.limit
            local req,_ = limit:get(token)
            if req then
                if req > CCcount then
                     log('GET',ngx.var.uri,"-","Url_List_Ip_CC")
                     ngx.exit(403)
                    return true
                else
                     limit:incr(token,1)
                end
            else
                limit:set(token,1,CCseconds)
            end
        end
    end
    return false
end

function base_ip_white_list_limit_module()
-- This module limit ip_white_list request rate ,if high concurrence ,request delay random time 
    if Ip_White_List_limit_Module then
        if ip_white_list_limit_rules ~="" and ngx.re.find(getClientIp(),ip_white_list_limit_rules,"imjo") then
            local uri=ngx.var.request_uri
            CCcount=tonumber(string.match(Ip_White_List_limit_rate,'(.*)/'))
            CCseconds=tonumber(string.match(Ip_White_List_limit_rate,'/(.*)'))
            local token = ip -- fingerprints 
            local limit = ngx.shared.limit
            local req,_=limit:get(token)
            if req then
                if req > CCcount then
                    log('GET',ngx.var.request_uri,"-","Ip_White_List_limit")
    		-- log('GET',ngx.var.request_uri,"-",ngx.now()*1000)
                    -- math.randomseed(ngx.now()*1000)
                    -- delay_time = math.random(1,tonumber(Ip_White_List_Delay_Max * 10))/10
    		-- log('GET',ngx.var.request_uri,"-",delay_time)
                    -- ngx.sleep(delay_time)
    		-- log('GET',ngx.var.request_uri,"-",delay_time)
    		ngx.exit(403)	
                    return true
                else
                     limit:incr(token,1)
                end
            else
                limit:set(token,1,CCseconds)
            end
        end
    end
    return false 
end 
       
function base_mime_type_and_ip_cc_module()
    if Mime_Type_Base_Ip_CC_Module then
        local current_mime_type = GetCurrentMimeType()
        if Mime_Type and ngx.re.find(current_mime_type, Mime_Type, "oi") then
            CCcount = tonumber(string.match(Mime_Type_Base_Ip_CC_Rate,'(.*)/'))
            CCseconds = tonumber(string.match(Mime_Type_Base_Ip_CC_Rate,'/(.*)'))
            local token = getClientIp()..current_mime_type -- fingerprints
            local limit = ngx.shared.limit
            local req,_ = limit:get(token)
            if req then
                if req > CCcount then
                    log('GET',ngx.var.request_uri,current_mime_type,Mime_Type)
                    ngx.exit(403)
                    return true
                else
                    limit:incr(token,1)
                end
            else
                limit:set(token,1,CCseconds)
            end
        end
    end
    return false
end 
